#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：yolo图像识别完整流程.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/7 上午9:22 
"""

import os
import shutil
import subprocess
from typing import List
import random
import json
import torch


class YOLOv5Pipeline:
	def __init__(self, input_dir: str, output_dir: str, train_ratio: float = 0.8):
		"""
		初始化YOLOv5流程。

		:param input_dir: 包含images和labels文件夹的输入目录路径。
		:param output_dir: 输出YOLO格式数据的目录路径。
		:param train_ratio: 训练集和测试集的比例，默认为0.8。
		"""
		self.input_dir = os.path.abspath(input_dir)
		self.output_dir = os.path.abspath(output_dir)
		self.train_ratio = train_ratio
		self.images_dir = os.path.join(self.input_dir, 'images')
		self.labels_dir = os.path.join(self.input_dir, 'labels')
		self.train_images_dir = os.path.join(self.output_dir, 'train', 'images')
		self.train_labels_dir = os.path.join(self.output_dir, 'train', 'labels')
		self.test_images_dir = os.path.join(self.output_dir, 'test', 'images')
		self.test_labels_dir = os.path.join(self.output_dir, 'test', 'labels')
		os.makedirs(self.train_images_dir, exist_ok=True)
		os.makedirs(self.train_labels_dir, exist_ok=True)
		os.makedirs(self.test_images_dir, exist_ok=True)
		os.makedirs(self.test_labels_dir, exist_ok=True)

		self.classes = self._get_classes()
		if not self.classes:
			print("警告：未找到任何类别标签。")
		self._write_classes_file()

	def _get_classes(self) -> List[str]:
		"""
		从JSON文件中提取所有类别。

		:return: 类别列表。
		"""
		classes = set()
		for filename in os.listdir(self.labels_dir):
			if filename.endswith('.json'):
				try:
					with open(os.path.join(self.labels_dir, filename), 'r') as f:
						data = json.load(f)
						for shape in data.get('shapes', []):
							classes.add(shape['label'])
				except Exception as e:
					print(f"读取文件 {filename} 时出错: {e}")
		return sorted(list(classes))

	def _write_classes_file(self):
		"""
		将类别写入classes.txt文件。
		"""
		if self.classes:
			with open(os.path.join(self.output_dir, 'classes.txt'), 'w') as f:
				for cls in self.classes:
					f.write(cls + '\n')

	def _convert_shape(self, shape: dict, img_width: int, img_height: int) -> List[float]:
		"""
		将Labelme的矩形形状转换为YOLO格式。

		:param shape: Labelme形状字典。
		:param img_width: 图像宽度。
		:param img_height: 图像高度。
		:return: YOLO格式的中心点坐标和宽高。
		"""
		x1, y1 = shape['points'][0]
		x2, y2 = shape['points'][1]
		x_center = (x1 + x2) / 2 / img_width
		y_center = (y1 + y2) / 2 / img_height
		width = (x2 - x1) / img_width
		height = (y2 - y1) / img_height
		return [x_center, y_center, width, height]

	def _convert_json_to_yolo(self, json_path: str, output_dir: str):
		"""
		将单个JSON文件转换为YOLO格式。

		:param json_path: JSON文件路径。
		:param output_dir: 输出目录。
		"""
		with open(json_path, 'r') as f:
			data = json.load(f)
			img_width, img_height = data['imageWidth'], data['imageHeight']
			txt_filename = os.path.splitext(os.path.basename(json_path))[0] + '.txt'
			with open(os.path.join(output_dir, txt_filename), 'w') as txt_file:
				for shape in data.get('shapes', []):
					if shape['shape_type'] == 'rectangle':
						x_center, y_center, width, height = self._convert_shape(shape, img_width, img_height)
						class_id = self.classes.index(shape['label'])
						txt_file.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

	def _copy_image(self, image_path: str, output_dir: str):
		"""
		复制图像文件到指定目录。

		:param image_path: 图像文件路径。
		:param output_dir: 输出目录。
		"""
		try:
			shutil.copy(image_path, output_dir)
		except FileNotFoundError:
			print(f"图像文件 {image_path} 未找到，跳过复制。")
		except Exception as e:
			print(f"复制图像文件 {image_path} 时出错: {e}")

	def prepare_dataset(self):
		"""
		准备数据集，包括转换标签和复制图像。
		"""
		json_files = [f for f in os.listdir(self.labels_dir) if f.endswith('.json')]
		random.shuffle(json_files)
		train_size = int(len(json_files) * self.train_ratio)
		train_files = json_files[:train_size]
		test_files = json_files[train_size:]

		for json_file in train_files:
			json_path = os.path.join(self.labels_dir, json_file)
			self._convert_json_to_yolo(json_path, self.train_labels_dir)
			image_filename = os.path.splitext(json_file)[0] + '.jpg'
			image_path = os.path.join(self.images_dir, image_filename)
			self._copy_image(image_path, self.train_images_dir)

		for json_file in test_files:
			json_path = os.path.join(self.labels_dir, json_file)
			self._convert_json_to_yolo(json_path, self.test_labels_dir)
			image_filename = os.path.splitext(json_file)[0] + '.jpg'
			image_path = os.path.join(self.images_dir, image_filename)
			self._copy_image(image_path, self.test_images_dir)

		# 检查目录是否为空
		if not any(os.scandir(self.train_images_dir)):
			print(f"警告：{self.train_images_dir} 目录为空，请检查数据准备过程。")
		if not any(os.scandir(self.train_labels_dir)):
			print(f"警告：{self.train_labels_dir} 目录为空，请检查数据准备过程。")
		if not any(os.scandir(self.test_images_dir)):
			print(f"警告：{self.test_images_dir} 目录为空，请检查数据准备过程。")
		if not any(os.scandir(self.test_labels_dir)):
			print(f"警告：{self.test_labels_dir} 目录为空，请检查数据准备过程。")

		print(f"数据集准备完成。训练集: {len(train_files)}个文件，测试集: {len(test_files)}个文件。")

	def train_model(self):
		"""
		训练YOLOv5模型。
		"""
		data_yaml_path = os.path.join(self.output_dir, 'data.yaml')
		weights_path = 'yolov5s.pt'  # 使用预训练的YOLOv5s模型
		output_dir = os.path.join(self.output_dir, 'runs', 'train')
		img_size = 640  # 图像大小
		batch_size = 8  # 批量大小
		epochs = 5  # 训练轮数
		device = 'cpu'  # 使用CPU
		workers = 2  # 数据加载器的工作线程数

		# 创建data.yaml文件
		with open(data_yaml_path, 'w') as f:
			f.write(f"train: {self.train_images_dir}\n")
			f.write(f"val: {self.test_images_dir}\n")
			f.write(f"nc: {len(self.classes)}\n")
			f.write(f"names: {self.classes}\n")

		# 检查data.yaml内容
		with open(data_yaml_path, 'r') as f:
			print("data.yaml 内容:")
			print(f.read())

		# 训练命令
		yolov5_dir = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本目录
		train_command = [
			"python", os.path.join(yolov5_dir, 'yolov5', 'train.py'),
			"--img", str(img_size),
			"--batch", str(batch_size),
			"--epochs", str(epochs),
			"--data", data_yaml_path,
			"--cfg", "yolov5/models/yolov5s.yaml",
			"--weights", weights_path,
			"--device", device,
			"--project", output_dir,
			"--workers", str(workers)
		]
		print("训练命令:", " ".join(train_command))
		result = subprocess.run(train_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
		print("Standard Output:", result.stdout.decode())
		print("Standard Error:", result.stderr.decode())
		if result.returncode != 0:
			print("训练失败，返回码:", result.returncode)
		else:
			print("训练成功")

	def test_model(self):
		"""
		测试YOLOv5模型。
		"""
		weights_path = os.path.join(self.output_dir, 'runs', 'train', 'exp', 'weights', 'best.pt')
		if not os.path.exists(weights_path):
			print(f"错误：模型权重文件 {weights_path} 不存在，请检查训练过程。")
			return

		test_dir = self.test_images_dir
		output_dir = os.path.join(self.output_dir, 'runs', 'test')
		img_size = 640  # 图像大小
		device = 'cpu'  # 使用CPU

		# 测试命令
		yolov5_dir = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本目录
		test_command = [
			"python", os.path.join(yolov5_dir, 'yolov5', 'detect.py'),
			"--weights", weights_path,
			"--source", test_dir,
			"--img", str(img_size),
			"--device", device,
			"--project", output_dir
		]
		print("测试命令:", " ".join(test_command))
		result = subprocess.run(test_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
		print("Standard Output:", result.stdout.decode())
		print("Standard Error:", result.stderr.decode())
		if result.returncode != 0:
			print("测试失败，返回码:", result.returncode)
		else:
			print("测试成功")


# 使用示例
if __name__ == "__main__":
	input_dir = 'datasets/HELMET/source_data'
	output_dir = 'datasets/HELMET/yolo_output_dir'
	pipeline = YOLOv5Pipeline(input_dir, output_dir)
	pipeline.prepare_dataset()
	pipeline.train_model()
	pipeline.test_model()
